# GDP与用电量关联分析散点图.py
# 功能：读取GDP和用电量累计值，计算同比增速，绘制散点图（X轴：用电量同比，Y轴：GDP同比，下拉切换省份）

import os
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots


def save_data_dic_to_excel(data_dic, output_path):
    """
    将data_dic中的数据按省份分sheet保存到Excel文件
    
    参数:
    data_dic: 字典，键为省份名称，值为包含'power'和'gdp'列的DataFrame
    output_path: 输出Excel文件的完整路径
    """
    # 创建ExcelWriter对象
    with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
        # 遍历字典中的每个省份数据
        for province, df in data_dic.items():
            # 将省份数据写入对应的sheet，sheet名称为省份名
            df.to_excel(writer, sheet_name=province)
        
        # 可选：创建一个汇总sheet，记录所有省份名称
        summary_df = pd.DataFrame({'省份': list(data_dic.keys())})
        summary_df.to_excel(writer, sheet_name='省份汇总', index=False)
    
    print(f"数据已成功保存到 {output_path}，共 {len(data_dic)} 个省份sheet")


# ------------------ 路径设置 ------------------
abs_path = os.path.abspath('.')
if '关联分析' in abs_path:
    abs_path = abs_path.replace('关联分析', '')

data_dir = os.path.join(f'{abs_path}/关联分析', 'data')
gdp_path = os.path.join(data_dir, 'GDP累计值.xlsx')
power_path = os.path.join(data_dir, '用电量累计值.xlsx')
output_html = os.path.join(f'{abs_path}/output', 'GDP与用电量同比增速散点图分析.html')
output_excel = os.path.join(f'{abs_path}/output', 'GDP与用电量同比增速数据.xlsx')


# ------------------ 读取数据 ------------------
# GDP：跳过第一行，第二行为列名，第一列为日期
gdp_df = pd.read_excel(gdp_path, header=1, index_col=0)
gdp_df.index = pd.to_datetime(gdp_df.index)

# 用电量：第一行为列名，第一列为日期
power_df = pd.read_excel(power_path, header=0, index_col=0)
power_df.index = pd.to_datetime(power_df.index)

# ------------------ 提取省份 ------------------
def extract_province(col):
    return col.split(':')[0]

gdp_cols = {extract_province(col): col for col in gdp_df.columns}
power_cols = {extract_province(col): col for col in power_df.columns}

# 共有省份
common_provinces = sorted(set(gdp_cols.keys()) & set(power_cols.keys()))

if len(common_provinces) == 0:
    raise ValueError("没有共同的省市数据，请检查列名格式是否为 '省份:指标:类型'")

print(f"共找到 {len(common_provinces)} 个共有省市：{common_provinces}")

# ------------------ 计算同比函数 ------------------
def calculate_yoy(series):
    """计算同比增速：(当前值 - 去年同期值) / |去年同期值|，返回单位为%的Series"""
    series = series.dropna()
    yoy = []
    dates = []
    for date, value in series.items():
        if date.month in [3, 6, 9, 12]:  # 只处理季度数据
            last_year_date = date - pd.DateOffset(years=1)
            if last_year_date in series.index:
                last_year_value = series[last_year_date]
                if abs(last_year_value) > 0:  # 避免除以0
                    growth = (value - last_year_value) / abs(last_year_value) * 100
                    yoy.append(growth)
                    dates.append(date)
    return pd.Series(yoy, index=dates).sort_index()

data_dic = {}

for province in common_provinces:
    gdp_col = gdp_cols[province]
    power_col = power_cols[province]

    # 计算同比并合并数据
    gdp_yoy = calculate_yoy(gdp_df[gdp_col].dropna())
    power_yoy = calculate_yoy(power_df[power_col].dropna())
    merged_data = pd.concat([power_yoy, gdp_yoy], axis=1, keys=['power', 'gdp']).dropna()
    data_dic[province] = merged_data


def plot_gdp_power_scatter_with_arrows(df, title="GDP与用电量关联分析"):
    """
    绘制GDP与用电量的散点图，并添加表示时间序列方向的箭头
    
    参数:
    df: DataFrame, 索引为日期格式，包含'power'和'gdp'两列
    title: 图表标题
    
    返回:
    fig: plotly图形对象
    """
    
    # 确保索引是日期格式
    df = df.copy()
    df.index = pd.to_datetime(df.index)
    
    # 按日期排序
    df = df.sort_index()
    
    # 创建图表
    fig = go.Figure()
    
    for i in range(len(df) - 1):
        ddf = df.iloc[i:i+2]
        from_x = ddf['power'].iloc[0]
        from_y = ddf['gdp'].iloc[0]
        to_x = ddf['power'].iloc[1]
        to_y = ddf['gdp'].iloc[1]

        if (to_x - from_x) * (to_y - from_y) > 0:
            color = 'green'
        else:
            color = 'red'

        # 添加散点和连线
        fig.add_trace(go.Scatter(
            x=ddf['power'],
            y=ddf['gdp'],
            mode='lines+markers',

            name='数据点',
            marker=dict(
                size=8,
                color='blue',
                opacity=0.7
            ),
            line=dict(
                color=color,
                width=2,
                dash='dot'
            ),
            text=[f"日期: {idx.strftime('%Y-%m')}<br>用电量: {x:.2f}<br>GDP: {y:.2f}" 
                for idx, x, y in zip(df.index, ddf['power'], ddf['gdp'])],
            hoverinfo='text'
        ))
    
    # 添加箭头注释
    annotations = []
    for i in range(len(df) - 1):
        t = df.index[i+1].strftime('%Y-%m')

        # 从当前点到下一个点添加箭头
        annotations.append(dict(
            x=df['power'].iloc[i+1],
            y=df['gdp'].iloc[i+1],
            xref='x',
            yref='y',
            showarrow=True,
            arrowhead=1,
            arrowcolor='red',
            arrowsize=1,
            arrowwidth=1,
            ax=- (df['power'].iloc[i+1] - df['power'].iloc[i]),
            ay=- (df['gdp'].iloc[i+1] - df['gdp'].iloc[i])
        ))

    annotations.append(dict(
            x=0.10,
            y=0.10,
            xref='x',
            yref='y',
            showarrow=True,
            arrowhead=2,
            arrowcolor='green',

            arrowsize=1,
            arrowwidth=1,
            ax=0.1,
            ay=0.1
        ))
        
    # 设置布局
    fig.update_layout(
        title=title,
        xaxis_title="用电量",
        yaxis_title="GDP",
        hovermode="closest",
        height=600,
        annotations=annotations,
        xaxis=dict(
            zeroline=True,
            zerolinecolor='gray',
            zerolinewidth=1
        ),
        yaxis=dict(
            zeroline=True,
            zerolinecolor='gray',
            zerolinewidth=1
        )
    )
    
    return fig

save_data_dic_to_excel(data_dic, output_excel)

# fg = plot_gdp_power_scatter_with_arrows(dic['上海'])

# 使用示例
# 假设您有一个DataFrame df，索引为日期，包含'power'和'gdp'列
# fig = plot_gdp_power_scatter_with_arrows(df)
# fig.show()  # 在Jupyter中显示
# 或者保存为HTML文件
# fig.write_html("scatter_plot_with_arrows.html")

